# -*- coding: utf-8 -*-
"""
@Time ： 2023-05-29 15:51
@Author ： Jinbo CHEN
@File ：dwriter.py
"""
import random

from adjustText import adjust_text
import math
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
import xlwings as xw
import os


def write_grey_matrices(excel_path, greys):
    mode = 'a' if os.path.isfile(excel_path) else 'w'
    with pd.ExcelWriter(excel_path, engine='openpyxl', mode=mode) as w:
        # grey matrices output
        matrix_names = ['df_lower', 'df_upper', 'norm_lower', 'norm_upper', 'Y', 'Z']
        for name in matrix_names:
            r_no = 0
            for g in greys:
                df = getattr(g, name)
                df.to_excel(w, sheet_name=name, startrow=r_no, startcol=0, index=True)
                r_no += len(g.df_lower)
                r_no += 2
    print("All grey matrix written into Excel!")
    return 0


def write_dematel_matrices(excel_path, dematel):
    mode = 'a' if os.path.isfile(excel_path) else 'w'
    with pd.ExcelWriter(excel_path, engine='openpyxl', mode=mode) as w:
        # dematel matrices output
        dematel.Z.to_excel(w, sheet_name='overall_Z', index=True)
        dematel.N.to_excel(w, sheet_name='overall_N', index=True)
        dematel.T.to_excel(w, sheet_name='overall_T', index=True)
        dematel.RDPE.to_excel(w, sheet_name='RDPE', index=True)

    with pd.ExcelWriter(excel_path, engine='openpyxl', mode='a', if_sheet_exists='overlay') as w:
        st_r = len(dematel.RDPE) + 2
        for rdpe in dematel.RDPEs:
            rdpe.to_excel(w, sheet_name='RDPE', startrow=st_r, index=True)
            st_r += (len(rdpe) + 2)
    print("All DEMATEL matrix written into Excel!")
    return 0


def write_danp_matrix(excel_path, danp):
    mode = 'a' if os.path.isfile(excel_path) else 'w'
    with pd.ExcelWriter(excel_path, engine='openpyxl', mode=mode) as w:
        danp.T2_NORM.to_excel(w, sheet_name='T2_NORM', index='True')
        danp.T_WEIGHT.to_excel(w, sheet_name='T_WEIGHT', index='True')
        danp.USM.to_excel(w, sheet_name='USM', index='True')
        danp.WSM.to_excel(w, sheet_name='WSM', index='True')
        danp.LWSM.to_excel(w, sheet_name='LWSM', index='True')
        danp.LWSM1.to_excel(w, sheet_name='LWSM1', index='True')
    print("All DANP matrix written into Excel!")
    return 0


# causal relationship quadrant diagram
def draw_causal_diagram(excel_path: str, sheet_name: str, place_at: str, df_rdpe: pd.DataFrame, df_T: pd.DataFrame):
    sns.set(style="white")
    ax = sns.scatterplot(x="P", y="E", data=df_rdpe, marker='o', alpha=0.7)
    x_min, x_max = df_rdpe['P'].min(), df_rdpe['P'].max()
    y_min, y_max = df_rdpe['E'].min(), df_rdpe['E'].max()
    ax.set_xlim(x_min - 0.2, x_max + 0.2)
    ax.set_ylim(y_min - 0.2, y_max + 0.2)
    plt.axvline(x=df_rdpe['P'].mean(), linewidth=0.9, color='r', linestyle='dotted')  # middle dot line vertical
    plt.axhline(y=0, linewidth=0.9, color='r', linestyle='dotted')  # middle dot line horizontal
    plt.xlabel('Prominence (D + R)')
    plt.ylabel('Relation (D - R)')
    # set chinese font family
    plt.rcParams["font.sans-serif"] = ["SimHei"]
    plt.rcParams["axes.unicode_minus"] = False
    # mark texts to points
    texts = []
    for i, row in df_rdpe.iterrows():
        texts.append(plt.text(row['P'], row['E'], i, fontsize=20, ha='right', va='top'))
        # draw arrows
        targets = df_T.index.tolist().copy()
        targets.remove(i)
        for t in targets:
            if df_T.loc[i, t] > df_T.loc[t, i]:
                dx = df_rdpe.loc[t, 'P'] - row['P']
                dy = df_rdpe.loc[t, 'E'] - row['E']
                length = math.sqrt((y_max-y_min) ** 2 + (x_max - x_min) ** 2)
                # random RGB color for arrow line
                red = random.random()
                green = random.random()
                blue = random.random()
                plt.arrow(row['P'], row['E'], dx * 0.90, dy * 0.90,
                          head_width=length * 0.01, head_length=length * 0.02, ec=(red, green, blue), linestyle='dashed')
    adjust_text(texts)
    # plt.show()

    # draw the figure into Excel sheet
    app = xw.App(visible=False, add_book=False)
    wb = app.books.open(excel_path)
    sheet = wb.sheets[sheet_name]  # 选择工作表
    sheet.pictures.add(plt.gcf(), left=sheet.range(place_at).left, top=sheet.range(place_at).top)  # 插入图表
    plt.close()
    wb.save()
    wb.close()
    app.quit()
    print("quadrant plot is drawn!")
    return 0
